import os
import functools

SCRIPT_DIR: str = os.path.dirname(os.path.realpath(__file__))


def to_snake_case(name: str) -> str:
    return functools.reduce(lambda x, y: x + ('_' if y.isupper() else '') + y, name).lower()


def find_enum(lines) -> (int, str):
    import re
    regex = re.compile(f"enum (struct|class) (\\w*)Op")
    for i, line in enumerate(lines):
        if matches := regex.match(line):
            enum_name = matches.group(2)
            print(f"Found enum {enum_name}")
            return i, enum_name
    return -1, ""


def process():
    # get current script directory
    src_index = SCRIPT_DIR.rfind("src")
    assert src_index != -1, "src not found in the path"
    xir_dir = os.path.join(SCRIPT_DIR[:src_index], "include", "luisa", "xir")
    header_path = os.path.join(xir_dir, "op.h")
    enum_names = []
    with open(header_path, "r") as f:
        lines = [x for line in f.readlines() if (x := line.strip()) and not x.startswith("//")]
    with open(f"{SCRIPT_DIR}/op.cpp", "w") as cpp_file:
        cpp_file.write("// This file is generated by update_op_name_map.py\n")
        cpp_file.write("// Do not edit it manually.\n\n")
        cpp_file.write("#include <luisa/core/logging.h>\n")
        cpp_file.write("#include <luisa/core/stl/unordered_map.h>\n")
        cpp_file.write("#include <luisa/xir/op.h>\n")
        cpp_file.write("\n")
        cpp_file.write("namespace luisa::compute::xir {\n\n")
        # search for enum (struct|class) {op_name}Op in lines
        while True:
            enum_index, enum_name = find_enum(lines)
            if enum_index == -1:
                break
            enum_names.append(enum_name)
            lines = lines[enum_index + 1:]
            enum_end = lines.index("};")
            ops = [line.split(",")[0].strip() for line in lines[:enum_end]]
            lines = lines[enum_end + 1:]
            cpp_file.write(f"""luisa::string_view to_string({enum_name}Op op) noexcept {{
    using namespace std::string_view_literals;
    switch (op) {{
""")
            cpp_file.writelines([f"        case {enum_name}Op::{op}: return \"{op.lower()}\"sv;\n" for op in ops])
            cpp_file.write(f"""    }}
    LUISA_ERROR_WITH_LOCATION("Unknown {to_snake_case(enum_name)} operation (code = {{}}).", static_cast<uint32_t>(op));
}}

{enum_name}Op {to_snake_case(enum_name)}_op_from_string(luisa::string_view name) noexcept {{
    using namespace std::string_view_literals;
    static const luisa::unordered_map<luisa::string_view, {enum_name}Op> m{{
""")
            cpp_file.writelines([f"        {{\"{op.lower()}\"sv, {enum_name}Op::{op}}},\n" for op in ops])
            cpp_file.write(f"""    }};
    auto iter = m.find(name);
    LUISA_ASSERT(iter != m.end(), "Unknown {to_snake_case(enum_name)} operation: {{}}.", name);
    return iter->second;
}}

""")
        cpp_file.write("}// namespace luisa::compute::xir\n")
    with open(os.path.join(xir_dir, "op_to_string.inl.h"), "w") as inl_file:
        inl_file.write("// This file is generated by update_op_name_map.py\n")
        inl_file.write("// Do not edit it manually.\n\n")
        for enum_name in enum_names:
            inl_file.write(f"[[nodiscard]] LUISA_XIR_API luisa::string_view to_string({enum_name}Op op) noexcept;\n")
            inl_file.write(
                f"[[nodiscard]] LUISA_XIR_API {enum_name}Op {to_snake_case(enum_name)}_op_from_string(luisa::string_view name) noexcept;\n\n")


if __name__ == "__main__":
    process()
